#ifndef _BSPLINE_OPTIMIZER_H_
#define _BSPLINE_OPTIMIZER_H_

#include <Eigen/Eigen>
#include "bspline_opt/lbfgs.hpp"
#include "bspline_opt/colors.h"
#include <traj_utils/uniform_bspline.h>
#include <traj_utils/plan_container.hpp>
#include <path_searching/dyn_a_star.h>
#include <plan_env/entropy_map.h>
#include <ros/ros.h>

namespace ego_planner {

class ControlPoints
{
public:
  int size_;
  Eigen::MatrixXd points;
  // The first free point at the directions vector
  std::vector<std::vector<Eigen::Vector3d>> free_pts;
  // Direction vector, must be normalized.
  std::vector<std::vector<Eigen::Vector3d>> directions;

  // A flag that used in many places. Initialize it everytime before using it.
  std::vector<bool> flag_temp;

  void resize(const int size)
  {
    size_ = size;

    free_pts.clear();
    directions.clear();
    flag_temp.clear();

    points.resize(3, size);
    free_pts.resize(size);
    directions.resize(size);
    flag_temp.resize(size, false);
  }

  void segment(ControlPoints &buf, const int start, const int end)
  {
    if (start < 0 || end >= size_ || points.rows() != 3)
    {
      ROS_ERROR("Wrong segment index! start=%d, end=%d, size=%d,points=%d",
                start, end, size_, points.rows());
      return;
    }

    buf.resize(end - start + 1);
    buf.points = points.block(0, start, 3, end - start + 1);
    buf.size_ = end - start + 1;
    for (int i = start; i <= end; i++)
    {
      buf.free_pts[i - start] = free_pts[i];
      buf.directions[i - start] = directions[i];
    }
  }
};

class BsplineOptimizer
{

public:
  BsplineOptimizer() {}
  ~BsplineOptimizer() {}

  /* main API */
  void setEnvironment(const EntropyMap::Ptr &map);
  void setParam(ros::NodeHandle &nh);

  /* helper function */

  // required inputs
  void setControlPoints(const Eigen::MatrixXd &points) { cps_.points = points; }
  void setBsplineInterval(const double &ts) { bspline_interval_ = ts; };
  void setObsPoses(ObsPosesPtr obs_poses) { obs_poses_ = obs_poses; }
  void setObsTrajs(ObsTrajsPtr obs_trajs) { obs_trajs_ = obs_trajs; }
  void setPosControlPointsAndTs(Eigen::MatrixXd &ctrl_pts, double ts);
  void setLocalTargetPt(const Eigen::Vector3d &target) { local_target_pt_ = target; };
  void setKeyViewPoint(Eigen::Vector3d pt, ros::Time t)
  {
    key_view_pt_ = pt;
    t_key_view_ = t;
  }

  AStar::Ptr Astar_;
  std::vector<Eigen::Vector3d> ref_pts_;

  bool InitAndProcess(const Eigen::MatrixXd &pts_init);
  std::vector<ControlPoints> distinctiveTrajs();
  bool OptimizeTraj(Eigen::MatrixXd &optimal_points, double ts);
  bool OptimizeTraj(Eigen::MatrixXd &optimal_points, double &final_cost, const ControlPoints &control_points, double ts);
  bool OptimizeTrajRefine(const Eigen::MatrixXd &init_points, const double ts, Eigen::MatrixXd &optimal_points);

  double getSwarmClearance() { return swarm_clearance_; }

  vector<int> occlusion_index;
  vector<Eigen::Vector3d> occlusion_target_point, occlusion_gradient;

  vector<Eigen::Vector3d> vis_pk, vis_pk_grad, vis_pk_grad_real;

private:
  EntropyMap::Ptr grid_map_;
  double map_reso_;
  std::vector<std::pair<int, int>> occ_segs_; // the start and end ID of segments in obstacles
  ObsPosesPtr obs_poses_{nullptr};
  ObsTrajsPtr obs_trajs_{nullptr};

  enum FORCE_STOP_OPTIMIZE_TYPE {
    DONT_STOP,
    STOP_FOR_REBOUND,
    STOP_FOR_ERROR
  } force_stop_type_;

  // main input
  ControlPoints cps_;       // B-spline control points
  double bspline_interval_; // B-spline knot span
  Eigen::Vector3d end_pt_;  // end of the trajectory

  /* optimization parameters */
  int order_;                                      // bspline degree
  double lambda_smooth_;                           // jerk smoothness weight
  double lambda_collision_, new_lambda_collision_; // distance weight
  double lambda_feasible_;                         // feasibility weight
  double lambda_fitness_;                          // curve fitting
  double lambda_mov_obj_;                          // dynamic obs collision
  double lambda_key_view_;                         // key view point slow down

  double dist_safe_, swarm_clearance_; // safe distance
  double dist_safe_dyn_;               // safe distance for dynamic obs
  double dist_attention_;              // distance when dynamic obs cause attention
  double max_vel_, max_acc_;           // dynamic limits
  double max_vel_key_view_;            // max vel when approaching key view point
  double coff_slow_down_key_view_;     // slow down when approaching key view point

  int variable_num_; // optimization variables
  int iter_num_;     // iteration of the solver

  Eigen::Vector3d local_target_pt_;
  Eigen::Vector3d key_view_pt_;
  ros::Time t_key_view_;

  bool checkCollisionFine();
  bool checkCollisionCoarse();
  bool dragSegOutofObs(std::vector<std::pair<int, int>> segment_ids, bool recordSeg = false);

  /* cost function */
  /* calculate each part of cost function with control points q as input */
  // q contains all control points
  void calcSmoothnessCost(const Eigen::MatrixXd &q, double &cost, Eigen::MatrixXd &gradient, bool falg_use_jerk = true);
  void calcFeasibilityCost(const Eigen::MatrixXd &q, double &cost, Eigen::MatrixXd &gradient);
  void calcTerminalCost(const Eigen::MatrixXd &q, double &cost, Eigen::MatrixXd &gradient);
  void calcDistanceCostRebound(const Eigen::MatrixXd &q, double &cost, Eigen::MatrixXd &gradient, int iter_num, double smoothness_cost);
  void calcMovingObjCost(const Eigen::MatrixXd &q, double &cost, Eigen::MatrixXd &gradient);
  void calcFitnessCost(const Eigen::MatrixXd &q, double &cost, Eigen::MatrixXd &gradient);
  void calcESDFReboundCost(const Eigen::MatrixXd &q, double &cost, Eigen::MatrixXd &gradient);
  void calcKeyViewPtCost(const Eigen::MatrixXd &q, double &cost, Eigen::MatrixXd &gradient);

  static double costFunction(void *instance,
                             const Eigen::VectorXd &x,
                             Eigen::VectorXd &grad);
  static double costFunctionRefine(void *instance,
                                   const Eigen::VectorXd &x,
                                   Eigen::VectorXd &g);

  static int earlyExit(void *instance,
                       const Eigen::VectorXd &x,
                       const Eigen::VectorXd &g,
                       const double fx,
                       const double step,
                       const int k,
                       const int ls);

  bool optimize(double &final_cost);
  bool optimize_refine();
  void combineCost(const Eigen::VectorXd &x, Eigen::VectorXd &grad, double &f);
  void combineCostRefine(const Eigen::VectorXd &x, Eigen::VectorXd &grad, double &f);

  // TODO: visibility old code
  void calcVisibilityReboundCost(const Eigen::MatrixXd &q, double &cost, Eigen::MatrixXd &gradient);

  // dynamic obstacles
  vector<vector<int>> dyn_obs_cp_index_;
  vector<vector<Eigen::Vector3d>> dyn_obs_pos_;

  void saveAstarToTxt(const std::vector<Eigen::Vector3d> &vec, const std::string &filename);
  void saveConotrolPointsToTxt(const Eigen::MatrixXd &matrix, const std::string &filename);

public:
  typedef unique_ptr<BsplineOptimizer> Ptr;

  EIGEN_MAKE_ALIGNED_OPERATOR_NEW
};

} // namespace ego_planner
#endif